
import os,sys
osa = os.path.abspath
osd = os.path.dirname
cur_dir = osd(osa(__file__))
par_dir = osd(cur_dir)
sys.path.insert(0,par_dir)


import os
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default='4', help='CUDA device(s) to use')
parser.add_argument('-p', '--port', type=int, default=20022, help='opened port')
args, _ = parser.parse_known_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

from diffusers import FluxKontextPipeline
from util_flux import  resize_with_aspect
from MODEL_CKP import FLUX_KONTEXT
import torch

import gradio as gr

# 初始化pipeline
pipe = FluxKontextPipeline.from_pretrained(FLUX_KONTEXT, torch_dtype=torch.bfloat16)
# 阿里妈妈 lora 加速 8步 完成    
pipe.load_lora_weights("/data/models/FLUX.1-Turbo-Alpha")  
pipe.to("cuda")

def gradio_infer(img, prompt, guidance_scale, steps):
    if img is None or not prompt:
        return None
    img = img.convert('RGB')
    img = resize_with_aspect('', img_pil=img)
    with torch.no_grad():
        result = pipe(
            image=img,
            prompt=prompt,
            prompt_2=prompt,
            height=img.height,
            width=img.width,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
        ).images[0]
    return result

with gr.Blocks(title="Flux Kontext Generation") as demo:
    gr.Markdown("# Flux Kontext 生成演示")
    gr.Markdown("上传图片，输入正向提示词，调整参数，生成新图像。")

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="输入图片", type="pil", height=512)
            prompt = gr.Textbox(label="正向提示词 (prompt)", lines=3, placeholder="请输入正向提示词")
            guidance_scale = gr.Slider(1, 10, value=2.5, step=0.1, label="Guidance Scale")
            steps = gr.Slider(8, 50, value=8, step=1, label="推理步数 (steps)")
            run_btn = gr.Button("生成")
        with gr.Column():
            output_img = gr.Image(label="输出图片", type="pil", height=512)

    run_btn.click(
        gradio_infer,
        inputs=[input_img, prompt, guidance_scale, steps],
        outputs=output_img
    )

if __name__ == "__main__":
    demo.launch(server_name='0.0.0.0', server_port=args.port)
